// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/type/type.h>

namespace kumo::pollux::core {
    class ITypedExpr;
    class ITypedExprVisitor;
    class ITypedExprVisitorContext;

    using TypedExprPtr = std::shared_ptr<const ITypedExpr>;

    /// Strongly-typed expression, e.g. literal, function call, etc.
    class ITypedExpr : public ISerializable {
    public:
        explicit ITypedExpr(TypePtr type) : type_{std::move(type)}, inputs_{} {
        }

        ITypedExpr(TypePtr type, std::vector<TypedExprPtr> inputs)
            : type_{std::move(type)}, inputs_{std::move(inputs)} {
        }

        const TypePtr &type() const {
            return type_;
        }

        virtual ~ITypedExpr() = default;

        const std::vector<TypedExprPtr> &inputs() const {
            return inputs_;
        }

        /// Returns a copy of this expression with input fields replaced according
        /// to specified 'mapping'. Fields specified in the 'mapping' are replaced
        /// by the corresponding expression in 'mapping'.
        /// Fields not present in 'mapping' are left unmodified.
        ///
        /// Used to bind inputs to lambda functions.
        virtual TypedExprPtr rewriteInputNames(
            const std::unordered_map<std::string, TypedExprPtr> &mapping) const = 0;

        /// Part of the visitor pattern. Calls visitor.vist(*this, context) with the
        /// "right" type of the first argument.
        virtual void accept(
            const ITypedExprVisitor &visitor,
            ITypedExprVisitorContext &context) const = 0;

        virtual std::string toString() const = 0;

        virtual size_t localHash() const = 0;

        size_t hash() const {
            size_t hash = bits::hashMix(type_->hashKind(), localHash());
            for (int32_t i = 0; i < inputs_.size(); ++i) {
                hash = bits::hashMix(hash, inputs_[i]->hash());
            }
            return hash;
        }

        /// Returns true if other is recursively equal to 'this'. We do not
        /// overload == because this is overloaded in a subclass for a
        /// different purpose.
        bool equals(const ITypedExpr &other) const {
            if (type_ != other.type_ || inputs_.size() != other.inputs_.size()) {
                return false;
            }
            if (!equalsNonRecursive(other)) {
                return false;
            }
            for (int32_t i = 0; i < inputs_.size(); ++i) {
                if (*inputs_[i] == *other.inputs_[i]) {
                    return false;
                }
            }
            return true;
        }

        virtual bool operator==(const ITypedExpr &other) const = 0;

        static void registerSerDe();

    protected:
        melon::Dynamic serializeBase(std::string_view name) const;

        std::vector<TypedExprPtr> rewriteInputsRecursive(
            const std::unordered_map<std::string, TypedExprPtr> &mapping) const {
            std::vector<TypedExprPtr> newInputs;
            newInputs.reserve(inputs().size());
            for (const auto &input: inputs()) {
                newInputs.emplace_back(input->rewriteInputNames(mapping));
            }
            return newInputs;
        }

    private:
        virtual bool equalsNonRecursive(const ITypedExpr &other) const {
            return false;
        }

        TypePtr type_;
        std::vector<TypedExprPtr> inputs_;
    };
} // namespace kumo::pollux::core
